#!/usr/bin/env python
#
# count_cve.py
#
# A helper script to count the number of CVEs
# found in the specified Debian packages so far.
# This script checks CVE data in Debian security tracker
# by parsing a json file provided in the security tracker web site.
#
# Usage:
#   $ ./count_cve.py [pkg1 pkg2 ...]
# eg:
#   $ ./count_cve.py apt bash busybox openssl
#
# Copyright (c) 2019 TOSHIBA Corporation
#
# SPDX-License-Identifier: Apache-2.0
#

import sys
import urllib2
import json
import os
DEBIAN_SEC_TRACK_DATA_URL = 'https://security-tracker.debian.org/tracker/data/json'
DEBIAN_SEC_TRACK_PKG_URL_BASE = 'https://security-tracker.debian.org/tracker/source-package/'
LOCAL_CVE_DATA_FILE = "cve_data.json"


class Cve:
    def __init__(self):
        self.cve_data_json = dict()

    def load_cve_data(self):
        """
        Load cve data from local cve data base file "cve_data.json"
        if the local file not exist or the file is not latest then download from server
        :return: True if loaded successfully
        """
        try:
            data = urllib2.urlopen(DEBIAN_SEC_TRACK_DATA_URL)
            server_last_modified = data.headers['Last-Modified']
            update = True
            if os.path.exists(LOCAL_CVE_DATA_FILE):
                # check the last-modified-date saved in local file with the server last-modified-date
                with open(LOCAL_CVE_DATA_FILE, 'r') as f:
                    if f.readline().partition('=')[2].strip() == server_last_modified:
                        self.cve_data_json = json.loads(f.readline())
                        update = False
            if update:
                # save / update to local CVE data file with last modified time
                print("Updating CVE to latest data: " + server_last_modified)
                cve_data = data.read()
                self.cve_data_json = json.loads(cve_data)
                with open(LOCAL_CVE_DATA_FILE, 'w') as f:
                    f.write('Server-Last-Modified=' + server_last_modified + '\n')
                    f.write(cve_data)
        except Exception as e:
            print("Warning: Failed to save / update CVE data locally: " + str(e))
            return False
        return True

    def check_pkg(self, pkg_name):
        """
        Check if there is a source-package page to see if the name is correct
        :return True if there is a source-package page
        """
        try:
            response = urllib2.urlopen(DEBIAN_SEC_TRACK_PKG_URL_BASE + pkg_name)
            page = response.read()
            if page.find('Available versions') != -1:
                return True
        except Exception as e:
            print("ERROR: Failed to open a source-package page: " + str(e))
        return False

    def count_cve(self, pkg_name):
        """
        Searches CVE vulnerability count in the data
        :return: count (int)
        """
        if not self.cve_data_json or len(self.cve_data_json) == 0:
            print("WARN: CVE data is not loaded")
            return 0

        if pkg_name in self.cve_data_json:
            return len([1 for vuln in self.cve_data_json[pkg_name] if vuln.startswith('CVE-')])
        else:
            return 0


if __name__ == "__main__":

    cve = Cve()
    if not cve.load_cve_data():
        exit(1)

    print('-' * 40)
    print("Package Name = No.of vulnerabilities")
    print('-' * 40)
    for pkg in sys.argv[1:]:
        if cve.check_pkg(pkg):
            print(pkg + " = " + str(cve.count_cve(pkg)))
        else:
            print(pkg + " does not exist")
    print('-' * 40)
